#!/usr/bin/env python
# prepare_sample_data.py
import os
import random
import shutil

def ensure_dir(directory):
    """Ensure directory exists"""
    if not os.path.exists(directory):
        os.makedirs(directory)

def create_sample_data(source_lines=1000, dev_lines=100):
    """Create sample data files for demonstration purposes"""
    print("Creating sample data files...")
    
    # Create required directories
    data_dir = "./data"
    tm_dir = os.path.join(data_dir, "TM-training-set")
    dev_dir = os.path.join(data_dir, "Dev-set")
    ref_dir = os.path.join(data_dir, "Reference-for-evaluation")
    
    ensure_dir(tm_dir)
    ensure_dir(dev_dir)
    ensure_dir(ref_dir)
    
    # Sample Chinese sentences
    chinese_samples = [
        "我 喜欢 学习 中文 。",
        "今天 天气 很 好 。",
        "请 问 你 是 哪国 人 ？",
        "我 想 去 北京 旅游 。",
        "这个 问题 很 难 。",
        "我 不 会 说 中文 。",
        "他 是 我 的 朋友 。",
        "我们 明天 见 。",
        "这 本 书 很 有趣 。",
        "我 在 学校 学习 。"
    ]
    
    # Sample English translations
    english_samples = [
        "I like learning Chinese .",
        "The weather is nice today .",
        "May I ask which country are you from ?",
        "I want to travel to Beijing .",
        "This question is difficult .",
        "I cannot speak Chinese .",
        "He is my friend .",
        "See you tomorrow .",
        "This book is very interesting .",
        "I study at school ."
    ]
    
    # Create training files
    with open(os.path.join(tm_dir, "chinese.txt"), "w", encoding="utf-8") as f_zh:
        for _ in range(source_lines):
            f_zh.write(random.choice(chinese_samples) + "\n")
            
    with open(os.path.join(tm_dir, "english.txt"), "w", encoding="utf-8") as f_en:
        for _ in range(source_lines):
            f_en.write(random.choice(english_samples) + "\n")
    
    # Create dev set
    with open(os.path.join(dev_dir, "Niu.dev.txt"), "w", encoding="utf-8") as f_dev:
        for _ in range(dev_lines):
            f_dev.write(random.choice(chinese_samples) + "\n")
    
    # Create reference file
    with open(os.path.join(ref_dir, "Niu.test.reference"), "w", encoding="utf-8") as f_ref:
        for _ in range(dev_lines):
            f_ref.write(random.choice(english_samples) + "\n")
    
    print("Sample data created successfully!")
    print(f"- Training data: {source_lines} lines")
    print(f"- Development data: {dev_lines} lines")
    print("You can now run 'python preprocess.py' to build vocabulary.")

if __name__ == "__main__":
    create_sample_data() 